package org.example.sentinel;

import com.alibaba.csp.sentinel.Entry;
import com.alibaba.csp.sentinel.SphU;
import com.alibaba.csp.sentinel.slots.block.BlockException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

@Component
public class IpSentinelFilter implements WebFilter {
    
    @Autowired
    private IpMonitorService ipMonitorService;
    
    @Autowired
    private SentinelRuleManager ruleManager;

    
    private String getClientIp(ServerWebExchange exchange) {
        // 实现获取真实客户端IP的逻辑
        // 示例：从X-Forwarded-For头获取
        String xff = exchange.getRequest().getHeaders().getFirst("X-Forwarded-For");
        if (xff != null && !xff.isEmpty()) {
            return xff.split(",")[0].trim();
        }
        return exchange.getRequest().getRemoteAddress() != null ? 
               exchange.getRequest().getRemoteAddress().getAddress().getHostAddress() : "unknown";
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        String clientIp = getClientIp(exchange);

        // 检查IP是否在黑名单
        return ipMonitorService.isBlacklisted(clientIp)
                .flatMap(isBlacklisted -> {
                    if (isBlacklisted) {
                        // 添加流控规则（如果尚未添加）
                        return ruleManager.addRuleForIp(clientIp)
                                .then(Mono.defer(() -> {
                                    // 直接拒绝请求
                                    exchange.getResponse().setStatusCode(HttpStatus.FORBIDDEN);
                                    return exchange.getResponse().setComplete();
                                }));
                    }

                    // 使用Sentinel进行流控
                    try (Entry entry = SphU.entry("ipBasedControl")) {
                        return chain.filter(exchange)
                                .doOnError(throwable -> {
                                    // 记录异常IP
                                    ipMonitorService.recordException(clientIp).subscribe();
                                });
                    } catch (BlockException ex) {
                        // 被流控规则拦截
                        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                        return exchange.getResponse().setComplete();
                    }
                });
    }
}